op {
  graph_op_name: "XlaSplitND"
  visibility: HIDDEN
  in_arg {
    name: "input"
    description: <<END
Input tensor to split across all dimensions.
  }
  out_arg {
    name: "outputs"
    description: <<END
Output slices based on input and num_splits defined, in row-major order.
END
  }
  attr {
    name: "num_splits"
    description: <<END
Number of ways to split per dimension. Shape dimensions must be evenly
divisible.
END
  }
  attr {
    name: "paddings"
    description: <<END
Optional list of right paddings per dimension of input tensor to apply before
splitting. This can be used to make a dimension evenly divisible.
END
  }
  summary: "Splits input tensor across all dimensions."
  description: <<END
An op which slices the input tensor based on the given num_splits attribute,
pads slices optionally, and returned the slices. Slices are returned in
row-major order.

This op may be generated via the TPU bridge.

For example, with `input` tensor:
```
[[0, 1, 2],
 [3, 4, 5],
 [6, 7, 8]]
```
`num_splits`:
```
[2, 2]
```
and `paddings`:
```
[1, 1]
```
the expected `outputs` is:
```
[[0, 1],
 [3, 4]]
[[2, 0],
 [5, 0]]
[[6, 7],
 [0, 0]]
[[8, 0],
 [0, 0]]
```
END
}
